/*
 * Copyright (c) 2020, Yanhui Shi <lime.syh at gmail dot com>
 * All rights reserved.
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include "core.h"
#include "proc.h"

/* Total processes generated by libcsp plugin. */
extern size_t csp_procs_num;
/* Process size generated by libcsp plugin which is guaranteed to be 4k-bytes
 * alignment. */
extern size_t csp_procs_size[];

extern _Thread_local csp_core_t *csp_this_core;

#ifndef csp_with_sysmalloc
extern void *csp_mem_alloc(size_t pid, size_t size);
extern void csp_mem_free(size_t pid, void *obj);
#endif

csp_proc_t *csp_proc_new(int id, bool waited_by_parent) {
  csp_core_t *this_core = csp_this_core;
  size_t size = csp_procs_size[id];

#ifdef csp_with_sysmalloc
  uintptr_t base = (uintptr_t)malloc(size);
#else
  uintptr_t base = (uintptr_t)csp_mem_alloc(this_core->pid, size);
#endif

  if (base == (uintptr_t)NULL) {
    errno = ENOMEM;
    perror("libcsp failed to alloc new proc.");
    exit(EXIT_FAILURE);
  }

  csp_proc_t *proc = (csp_proc_t *)(base + size - sizeof(csp_proc_t));
  proc->base = base;
  proc->is_new = true;
  proc->borned_pid = this_core->pid;
  atomic_store(&proc->stat, csp_proc_stat_none);

  /* We should make sure %rbp is 16-bytes alignment. */
  proc->rbp = (uintptr_t)proc - (!!((uintptr_t)proc & 0x0f) << 3);

  proc->nchild = 0;
  if (waited_by_parent) {
    proc->parent = this_core->running;
  } else {
    proc->parent = NULL;
  }
  proc->pre = proc->next = NULL;

#ifdef csp_enable_valgrind
  proc->valgrind_stack = VALGRIND_STACK_REGISTER(proc->base, proc);
#endif
  return proc;
}

__attribute__((noinline,used)) void csp_proc_nchild_set(size_t nchild) {
  atomic_store(&csp_this_core->running->nchild, nchild);
}

__attribute__((naked)) void csp_proc_restore(csp_proc_t *proc) {
  __asm__ __volatile__(
    "ldmxcsr 0x18(%rdi)\n"
    "fldcw   0x1c(%rdi)\n"
    "mov     0x20(%rdi), %rsp\n"
    "mov     0x28(%rdi), %rbp\n"

    /* Load `is_new` and check it. */
    "mov  0x10(%rdi), %rax\n"
    "test %eax, %eax\n"
    "jne  1f\n"

    /* If it's 0, restore remain callee-saved registers. At most cases, we are
     * likely to yield many times in a process, so we put this code ahead to
     * achieve better performance according to the Branch Prediction model. */
    "mov 0x30(%rdi), %rbx\n"
    "mov 0x38(%rdi), %r12\n"
    "mov 0x40(%rdi), %r13\n"
    "mov 0x48(%rdi), %r14\n"
    "mov 0x50(%rdi), %r15\n"
    "retq\n"

    /* Otherwise, restore caller-saved registers. */
    "1: movq $0, 0x10(%rdi)\n" // Set `is_new` to 0.
    "mov 0x38(%rdi), %rsi\n"
    "mov 0x40(%rdi), %rdx\n"
    "mov 0x48(%rdi), %rcx\n"
    "mov 0x50(%rdi), %r8\n"
    "mov 0x58(%rdi), %r9\n"
    "mov 0x30(%rdi), %rdi\n" // Restore %rdi at the last step.
    "retq\n"
  );
}

__attribute__((noinline)) void csp_proc_destroy(csp_proc_t *proc) {
#ifdef csp_enable_valgrind
  VALGRIND_STACK_DEREGISTER(proc->valgrind_stack);
#endif

#ifdef csp_with_sysmalloc
  free((void *)proc->base);
#else
  csp_mem_free(proc->borned_pid, (void *)proc->base);
#endif
}
